package com.open.basic.trace;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

import com.open.basic.constant.TraceId;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.util.StringUtils;


/**
 * 任务装饰器
 */
@Slf4j
public class MdcTaskDecorator implements TaskDecorator {

    /**
     * 使异步线程池获得主线程的上下文
     * 装饰
     *
     * @param runnable
     * @return
     */
    @Override
    public Runnable decorate(Runnable runnable) {
        final Map<String, String> map = MDC.getCopyOfContextMap();
        return () -> {
            try {
                Map<String, String> traceMap = map == null ? new HashMap<>() : map;
                String traceId = traceMap.get(TraceId.TRACE_ID);
                if (!StringUtils.hasLength(traceId)) {
                    traceId = TraceId.getTraceId();
                }
                traceMap.put(TraceId.TRACE_ID, traceId);
                MDC.setContextMap(traceMap);
                runnable.run();
            } catch (Exception e) {
                log.error("traceId:{},task execute error:{}", TraceId.getTraceId(), e.getMessage(), e);
            } finally {
                MDC.clear();
            }
        };
    }
}